# coding=utf-8
# 学习单位  : 郑州大学
# @Author  : 铭同学
# @Time    : 2021/10/11 9:19
# @Software: PyCharm

import torch

# 准备数据
# x,y是3行1列的矩阵，也就是说共有三个样本数据（行），每个数据是一个特征（列）
x_data = torch.tensor([[1.0],[2.0],[3.0]])
y_data = torch.tensor([[2.0],[4.0],[6.0]])

# 涉及线性模型的类（模型的设计都要用类实现，并继承于torch.nn.model）
class LineraModel(torch.nn.Module):
    def __init__(self):
        super(LineraModel, self).__init__()
        # (1,1)是指输入x和输出y的特征维度，这里数据集中的x和y的特征都是1维的
        # 该线性层需要学习的参数是w和b  获取w/b的方式分别是~linear.weight/linear.bias
        self.linear = torch.nn.Linear(1,1)


    # 必须要定义forward同名函数，其实是对forward的实现
    def forward(self, x):
        y_pred = self.linear(x)
        return y_pred

# 模型的实例化
model = LineraModel()

# 定义损失和优化器
criterion = torch.nn.MSELoss(size_average=False, reduction='sum') # 这里不求均值，Ture就对误差求均值
optimizer = torch.optim.SGD(model.parameters(), lr=0.01) #  model.parameters()自动完成模型所有参数的初始化操作

# 开始训练轮次  forward、梯度清零、backward、update
for epoch in range(300):
    y_pred = model(x_data) # 预测--forward
    loss = criterion(y_pred, y_data) # 计算损失--forward
    print(f'训练轮次：{epoch}',f'本轮次损失：{loss.item()}')

    # the grad computer by .backward() will be accumulated. so before backward, remember set the grad to zero
    optimizer.zero_grad() # backward前梯度清零
    loss.backward() # backward自动计算梯度
    optimizer.step() # update 参数，对w（权重）和b（偏置）进行更新



print('w = ', model.linear.weight.item())
print('b = ', model.linear.bias.item())


# 用数据进行测试训练结果
x_test = torch.tensor([[4.0]])
y_test = model(x_test)
print(f'y_pred = {y_test.data}')



